package com.canyou.limiter;

import com.canyou.enums.ErrorCodeEnum;
import com.canyou.exception.CysBusinessException;
import com.canyou.utils.IpUtils;
import com.canyou.utils.StringUtils;
import com.google.common.collect.ImmutableList;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.aop.support.AopUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.io.Serializable;
import java.lang.reflect.Method;

/**
 * 限流切面
 *
 * @author fy
 * @date 20/07/10 17:06
 */
@Aspect
@Component
@Slf4j
public class RateLimiterAspect {

    @Autowired
    private RedisTemplate<String, Serializable> redisTemplate;

    @Around("execution(public * *(..)) && @annotation(com.canyou.limiter.RateLimiter)")
    public Object interceptor(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        RateLimiter limiterAnnotation = method.getAnnotation(RateLimiter.class);
        if (limiterAnnotation == null) {
            return pjp.proceed();
        }
        Class<?> targetClass = (pjp.getSignature() != null ? AopUtils.getTargetClass(pjp.getThis()) : null);
        LimiterTypeEnum limitType = limiterAnnotation.limitType();
        String prefix = String.format("rate-limiter%s#%s", targetClass != null ? targetClass.getName() : null, method.getName());
        String key;
        int limit = limiterAnnotation.count();
        int period = limiterAnnotation.period();

        switch (limitType) {
            case CUSTOMER:
                key = limiterAnnotation.key();
                if (StringUtils.isBlank(key)) {
                    key = prefix;
                }
                break;
            case IP:
            default:
                key = getIpAddress();
                break;
        }

        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(prefix, key));
        Number count = null;
        try {
            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            count = redisTemplate.execute(redisScript, keys, limit, period);
            log.info("Access try count is {} for key = {}", count, key);
        } catch (Exception e) {
            log.error("redis获取限流失败", e);
        }
        if (count != null && count.intValue() <= limit) {
            return pjp.proceed();
        } else {
            throw new CysBusinessException(ErrorCodeEnum.FREQUENT_REQUESTS, "您请求的太频繁了");
        }
    }

    /**
     * redis Lua 限流脚本
     *
     * @return
     */
    private String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 调用不超过最大值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }

    private String getIpAddress() {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        return IpUtils.getIpAddr(request);
    }

}
